#include "ShadowProcess.h"
#include "Utils/LogUtil.h"
#include "Render/RenderState.h"
#include "Core/CoreGlobals.h"
#include "Game/Scene/Entity.h"
#include "Game/Component/Model.h"
#include "Render/Buffer/VertexBuffer.h"
#include "Render/Buffer/IndexBuffer.h"
#include "Render/Buffer/CustomBuffer.h"
#include "LocalShader.h"

void ShadowProcess::OnInit()
{
	mDim = 1024;

	mModelEntitys = mGameScene->GetEntitysByTypes<LocalEntity, Model, TextureMaterial>();

	mLightEntitys = mGameScene->GetEntitysByType<LocalEntity, LightMaterial>();


	float fov = mGameScene->GetCamera().GetFovY();
	float aspect = mGameScene->GetCamera().GetAspect();
	float nearZ = mGameScene->GetCamera().GetNearZ();
	float farZ = mGameScene->GetCamera().GetFarZ();

	Ubo ubo;

	glm::mat4 shadowProj = glm::perspective(fov, aspect, nearZ, farZ);

	for (uint32 i = 0; i < 3; ++i)
	{
		LightMaterial& lightMaterial = mLightEntitys[i]->GetComponent<LightMaterial>();

		// ubo.mvp[i] = shadowProj * glm::lookAt(glm::vec3(mLightEntitys[i]->GetTransform().GetLocalPosition()), glm::vec3(lightMaterial.target), glm::vec3(0.0f, 1.0f, 0.0f));

		ubo.mvp[i] = mGameScene->GetCamera().GetProj() * mGameScene->GetCamera().GetView();
	}


	ENQUEUE_RENDER_COMMAND(BuildModel)([this, ubo](VkCommandBuffer commandBuffer)
		{
			Ubo uboData = ubo;
			mUboBuffer = CustomBuffer::Builder()
				.SetData(&uboData)
				.SetSize(sizeof(Ubo))
				.SetMemoryProperty(VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)
				.Build();

			mDescriptor = RHIDescriptor::Builder()
				.PushShader<ShadowVertexShader>()
				.PushShader<ShadowGeometryShader>()
				.SetDescriptorCount(0, mModelEntitys.size())
				.Build();

			mDescriptorMasks.resize(mModelEntitys.size());
			for (uint32 i = 0; i < mModelEntitys.size(); ++i)
			{
				mDescriptor->SetValue("vertUbo", mModelEntitys[i]->mUboBuffer);
				mDescriptor->SetValue("gemoUbo", mUboBuffer);
				mDescriptorMasks[i] = mDescriptor->AllocateDescriptor();
			}

			mShadowRenderTarget = DepthRenderTarget::Builder()
				.SetWidth(mDim)
				.SetHeight(mDim)
				.SetArrayLayers(3)
				.SetUsage(VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetAspectMask(VK_IMAGE_ASPECT_DEPTH_BIT)
				.SetFormat(VK_FORMAT_D32_SFLOAT_S8_UINT)
				.SetMultiple(true)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				// .SetClearDepthStencilValue({0.5f, 0})
				.Build(this, "ShadowRenderTarget");

			mShadowSampler = RHISampler::Builder()
				.SetAddressMode(VK_SAMPLER_ADDRESS_MODE_CLAMP_TO_EDGE)
				.Build(this, "ShadowSampler");

			mPipeline = GraphicsPipeline::Builder()
				.SetDescriptor(mDescriptor)
				.SetDepthRenderTarget(mShadowRenderTarget)
				.SetCullMode(VK_CULL_MODE_FRONT_BIT)
				.SetDepthCompareOp(VK_COMPARE_OP_LESS_OR_EQUAL)
				.SetDepthBiasEnable(VK_TRUE)
				.AddDynamicState(VK_DYNAMIC_STATE_DEPTH_BIAS)
				.SetDstStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetDstAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetSrcStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetSrcAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.Build();

			mRenderPass = RHIRenderPass::Builder()
				.AddRenderTarget(mShadowRenderTarget)
				.AddPipeline(mPipeline)
				.SetSrcStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetSrcAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetDstAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.Build();

			mFrameBuffer = RHIFrameBuffer::Builder()
				.SetRenderPass(mRenderPass)
				// Important, waste my all day
				.SetLayers(3)
				.Build();

			mCommandBuffer = GetCommandBuffer("MrtProcess", "OfflineCommandBuffer");

			mWaitGraphicsSemaphore = GetSemaphore("MrtProcess", "WaitGraphicsSemaphore");

			mSignalOfflineSemaphore = GetSemaphore("MrtProcess", "SignalOfflineSemaphore");
		});
}

void ShadowProcess::OnTick(float deltaTime)
{
	float fov = mGameScene->GetCamera().GetFovY();
	float aspect = mGameScene->GetCamera().GetAspect();
	float nearZ = mGameScene->GetCamera().GetNearZ();
	float farZ = mGameScene->GetCamera().GetFarZ();

	Ubo ubo;

	glm::mat4 shadowProj = glm::perspective(fov, aspect, nearZ, farZ);

	for (uint32 i = 0; i < 3; ++i)
	{
		LightMaterial& lightMaterial = mLightEntitys[i]->GetComponent<LightMaterial>();

		ubo.mvp[i] = shadowProj * glm::lookAt(glm::vec3(mLightEntitys[i]->GetTransform().GetLocalPosition()), glm::vec3(lightMaterial.target), glm::vec3(0.0f, -1.0f, 0.0f));

		// ubo.mvp[i] = mGameScene->GetCamera().GetProj() * mGameScene->GetCamera().GetView();
	}


	ENQUEUE_RENDER_COMMAND(RecordCommand)([this, ubo](VkCommandBuffer commandBuffer)
		{
			Ubo uboData = ubo;
			mUboBuffer->WriteMemory(&uboData);

			VkCommandBuffer offlineCommandBuffer = mCommandBuffer->GetCommandBuffer();

			VkRenderPassBeginInfo renderPassBeginInfo = mFrameBuffer->GetRenderPassBeginInfo();

			vkCmdBeginRenderPass(offlineCommandBuffer, &renderPassBeginInfo, VK_SUBPASS_CONTENTS_INLINE);


			VkViewport viewPort = { 0, 0, mDim, mDim, 0, 1 };
			VkRect2D scissor = {0, 0, mDim, mDim};
			vkCmdSetViewport(offlineCommandBuffer, 0, 1, &viewPort);
			vkCmdSetScissor(offlineCommandBuffer, 0, 1, &scissor);

			// Set depth bias (aka "Polygon offset")
			vkCmdSetDepthBias(
				offlineCommandBuffer,
				1.25f,
				0.0f,
				1.75f);

			vkCmdBindPipeline(offlineCommandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipeline());

			for (uint32 i = 0; i < mModelEntitys.size(); ++i)
			{
				vkCmdBindDescriptorSets(offlineCommandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipelineLayout(), 0, 1, &mDescriptorMasks[i].GetDescriptorSet(), 0, nullptr);

				VkBuffer vertexBuffer = mModelEntitys[i]->mVertexBuffer->GetBuffer();
				VkDeviceSize offsets = 0;
				vkCmdBindVertexBuffers(offlineCommandBuffer, 0, 1, &vertexBuffer, &offsets);

				VkBuffer indexBuffer = mModelEntitys[i]->mIndexBuffer->GetBuffer();
				vkCmdBindIndexBuffer(offlineCommandBuffer, indexBuffer, 0, VK_INDEX_TYPE_UINT32);

				vkCmdDrawIndexed(offlineCommandBuffer, mModelEntitys[i]->mIndexBuffer->GetIndexCount(), 1, 0, 0, 0);
			}

			vkCmdEndRenderPass(offlineCommandBuffer);

			mCommandBuffer->FlushCommandBuffer(offlineCommandBuffer, { mWaitGraphicsSemaphore }, { mSignalOfflineSemaphore });
		});
}

void ShadowProcess::OnResize()
{
	
}
